package com.ruyuan.rapid.core.balance;

import com.ruyuan.rapid.common.config.ServiceInstance;
import com.ruyuan.rapid.common.enums.LoadBalanceStrategy;
import com.ruyuan.rapid.common.enums.ResponseCode;
import com.ruyuan.rapid.common.exception.RapidResponseException;
import com.ruyuan.rapid.core.context.AttributeKey;
import com.ruyuan.rapid.core.context.RapidContext;
import com.ruyuan.rapid.core.helper.DubboReferenceHelper;
import org.apache.dubbo.common.URL;
import org.apache.dubbo.common.constants.CommonConstants;
import org.apache.dubbo.rpc.Invocation;
import org.apache.dubbo.rpc.Invoker;
import org.apache.dubbo.rpc.RpcContext;
import org.apache.dubbo.rpc.RpcException;
import org.apache.dubbo.rpc.cluster.Constants;

import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;

import static org.apache.dubbo.common.constants.CommonConstants.TIMESTAMP_KEY;

/**
 * Dubbo 负载均衡
 */
public class DubboLoadBalance implements org.apache.dubbo.rpc.cluster.LoadBalance {

    public static final String NAME = "rlb";

    /**
     * select one invoker in list.
     * @param invokers   invokers.
     * @param url        refer url
     * @param invocation invocation.
     * @return selected invoker.
     */
    @Override
    public <T> Invoker<T> select(List<Invoker<T>> invokers, URL url, Invocation invocation) throws RpcException {
        System.out.println("-----------------------Dubbo LoadBalance Into -----------------------");
        RapidContext rapidContext = (RapidContext) RpcContext.getContext().get(DubboReferenceHelper.DUBBO_TRANSFER_CONTEXT);
        LoadBalanceStrategy loadBalanceStrategy = rapidContext.getAttribute(AttributeKey.DUBBO_LOADBALANCE_STRATEGY);
        LoadBalance loadBalance = LoadBalanceFactory.getLoadBalance(loadBalanceStrategy);
        Set<ServiceInstance> instanceWrappers = new HashSet<>();
        for (Invoker<T> invoker : invokers) {
            instanceWrappers.add(new ServiceInstanceWrapper(invoker, invocation));
        }

        // 把dubbo invokers的服务列表实例,转成自己能够认识的ServcieInstance实例, 设置到全局上下文里
        rapidContext.putAttribute(AttributeKey.MATCH_INSTANCES, instanceWrappers);
        ServiceInstance serviceInstance = loadBalance.select(rapidContext);
        if (serviceInstance instanceof ServiceInstanceWrapper) {
            return ((ServiceInstanceWrapper) serviceInstance).getInvoker();
        } else {
            throw new RapidResponseException(ResponseCode.SERVICE_INSTANCE_NOT_FOUND);
        }
    }

    static class ServiceInstanceWrapper<T> extends ServiceInstance {

        private final Invoker<T> invoker;

        public ServiceInstanceWrapper(Invoker<T> invoker, Invocation invocation) {
            this.invoker = invoker;
            this.setAddress(invoker.getUrl().getAddress());
            this.setUniqueId(invoker.getUrl().getServiceKey());
            this.setRegisterTime(invoker.getUrl().getParameter(TIMESTAMP_KEY, 0L));
            this.setWeight(invoker.getUrl().getMethodParameter(invocation.getMethodName(), Constants.WEIGHT_KEY, Constants.DEFAULT_WEIGHT));
            this.setVersion(invoker.getUrl().getParameter(CommonConstants.VERSION_KEY));
            this.setEnable(true);
        }

        public Invoker<T> getInvoker() {
            return invoker;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) {
                return true;
            }
            if (o == null || getClass() != o.getClass()) {
                return false;
            }
            if (!super.equals(o)) {
                return false;
            }
            ServiceInstanceWrapper<?> that = (ServiceInstanceWrapper<?>) o;
            return Objects.equals(this.getAddress(), that.getAddress());
        }

        @Override
        public int hashCode() {
            return Objects.hash(this.getAddress());
        }
    }
}
